// smart_tensor.cpp
#include "smart_tensor.h"
#include <iostream>

// 辅助函数：将多维索引转换为一维索引
int SmartTensor::flattenIndex(const std::vector<int>& indices) const {
    int index = 0;
    int multiplier = 1;
    for (int i = shape.size() - 1; i >= 0; --i) {
        index += indices[i] * multiplier;
        multiplier *= shape[i];
    }
    return index;
}

// 构造函数
SmartTensor::SmartTensor(const std::vector<int>& s) : shape(s) {
    int size = 1;
    for (int dim : shape) {
        size *= dim;
    }
    data.resize(size, 0.0);
}

// 获取张量的形状
std::vector<int> SmartTensor::getShape() const {
    return shape;
}

// 获取元素
double SmartTensor::get(const std::vector<int>& indices) const {
    if (indices.size() != shape.size()) {
        throw std::invalid_argument("Indices size does not match tensor shape");
    }
    return data[flattenIndex(indices)];
}

// 设置元素
void SmartTensor::set(const std::vector<int>& indices, double value) {
    if (indices.size() != shape.size()) {
        throw std::invalid_argument("Indices size does not match tensor shape");
    }
    data[flattenIndex(indices)] = value;
}

// 张量加法
SmartTensor SmartTensor::operator+(const SmartTensor& other) const {
    if (shape != other.shape) {
        throw std::invalid_argument("Tensor shapes must match for addition");
    }
    SmartTensor result(shape);
    for (size_t i = 0; i < data.size(); ++i) {
        result.data[i] = data[i] + other.data[i];
    }
    return result;
}

// 标量乘法
SmartTensor SmartTensor::operator*(double scalar) const {
    SmartTensor result(shape);
    for (size_t i = 0; i < data.size(); ++i) {
        result.data[i] = data[i] * scalar;
    }
    return result;
}

// 矩阵乘法（仅适用于二阶张量）
SmartTensor SmartTensor::matrixMultiply(const SmartTensor& other) const {
    if (shape.size() != 2 || other.shape.size() != 2) {
        throw std::invalid_argument("Matrix multiplication requires second - order tensors");
    }
    if (shape[1] != other.shape[0]) {
        throw std::invalid_argument("Number of columns in the first matrix must be equal to the number of rows in the second matrix");
    }
    std::vector<int> resultShape = {shape[0], other.shape[1]};
    SmartTensor result(resultShape);
    for (int i = 0; i < shape[0]; ++i) {
        for (int j = 0; j < other.shape[1]; ++j) {
            double sum = 0.0;
            for (int k = 0; k < shape[1]; ++k) {
                std::vector<int> index1 = {i, k};
                std::vector<int> index2 = {k, j};
                sum += get(index1) * other.get(index2);
            }
            std::vector<int> resultIndex = {i, j};
            result.set(resultIndex, sum);
        }
    }
    return result;
}

// 打印张量
void SmartTensor::print() const {
    std::cout << "Shape: ";
    for (int dim : shape) {
        std::cout << dim << " ";
    }
    std::cout << std::endl;
    std::cout << "Data: ";
    for (double val : data) {
        std::cout << val << " ";
    }
    std::cout << std::endl;
}